#!/usr/bin/env python
#
'''
Module to provide a more pleasurable pythonic experience when working with
FITS files
'''
import sys

import ctypes
import cfitsio

from pyerr import *

# Pythonize the opening of FITS files
def open(filename, mode='r', clobber=False):
    '''Factory function to create a fitsfile object in a pythonic manner'''

    ff = FitsFile()

    # Open existing FITS file
    if mode=='r' or mode=='rw':
        if mode=='r':
            ff.open(filename, mode=cfitsio.READONLY)
        if mode=='rw':
            ff.open(filename, mode=cfitsio.READWRITE)

    # Open a new FITS for writing
    elif mode=='w':

        # Overwrite any existing file if clobber option is set
        if clobber:
            ff.create_file('!' + filename)

        # If not clobbering and the file already exists, a FitsioError
        # exception will be raised
        else:
            ff.create_file(filename)

    else:
        raise ValueError('Mode must be \'r\', \'rw\', or \'w\'')

    return ff

#------------------------------------------------------------------------------

# Some common and simple routines in a pythonic manner
#
# The point here is that we just use basic file level operations

def RUNCF(ff, funcname, *args):
    '''Generalized wrapper to most CFITSIO functions'''
    
    # Most CFITSIO functions are of the form:
    # functionname(fptr, ... arg list ..., status)
    # This just provides meta-code of a sort to call them
    
    # List of parameter arguments. need to be a list rather than a tuple so 
    # that we can insert other entries
    arglist = list(args)

    # Pre-pend the pointer to the fitsfile()
    arglist.insert(0, ff.fptr)

    # Append the reference to the status value
    ff.status = ctypes.c_int(0)
    arglist.append(ctypes.byref(ff.status))

    # Call the routine given by funcname and deal with the error in case 
    # of a non-zero return value
    if funcname(*arglist):
        handle_error(ff.status)

class FitsFile(object):
    '''Class to represent a FITS file.

    There are many different combinations of ways the CFITSIO functions can
    be used and thus many types of possible OO designs for a FITS file
    class. This class essentially just makes it easier to create and access
    the fitsfile pointer that is passed into the FITSIO routines.

    Attributes:
    fptr: a ctype pointer that maps to the C fitsfile pointer
    filename: the name of the FITS format file.
    '''

    def __init__(self):
        '''Initialize with an empty C fitsfile pointer'''
        
        self.fptr = ctypes.pointer( cfitsio.fitsfile() )
        self.status = ctypes.c_int(0)
        self.filename = None

        self.value = ctypes.create_string_buffer('\000' * cfitsio.FLEN_VALUE)
        self.comment = ctypes.create_string_buffer('\000' * cfitsio.FLEN_COMMENT)

    def open(self, filename, mode=cfitsio.READONLY):
        '''Wrapper to ffopen / fits_open_file'''
        
        self.filename = filename
        if cfitsio.ffopen(ctypes.byref(self.fptr), filename, mode, 
                          ctypes.byref(self.status)):
            handle_error(self.status)

    def create_file(self, filename):
        '''Wrapper to ffinit / fits_create_file'''
        
        self.filename = filename
        if cfitsio.ffinit(ctypes.byref(self.fptr), filename,
                          ctypes.byref(self.status)):
            handle_error(self.status)

    def create_table(self, tbltype, extname, ttype, tform, tunit, nrows):
        '''Wrapper to ffcrtb / fits_create_tbl'''
        
        # Number of columns (to do: check for compatible sizes)
        ncols = len(ttype)
        tfields = ctypes.c_int(ncols)
        naxis2 = ctypes.c_longlong(nrows)
        RUNCF(self, cfitsio.ffcrtb, tbltype, naxis2, tfields,
              (ctypes.c_char_p * ncols)(*ttype),
              (ctypes.c_char_p * ncols)(*tform),
              (ctypes.c_char_p * ncols)(*tunit),
              extname)

    def get_num_hdus(self):
        '''Wrapper to ffthdu / fits_get_num_hdus'''

        numhdus = ctypes.c_int()
        RUNCF(self, cfitsio.ffthdu, ctypes.byref(numhdus))
        return numhdus.value

    def move_to_hdu(self, hdunum):
        '''Wrapper to ffmahd / fits_movabs_hdu'''

        hdutype = ctypes.c_int()
        RUNCF(self, cfitsio.ffmahd, hdunum, hdutype)
        return hdutype.value

    def create_hdu(self):
        '''Wrapper to ffcrhd/fits_create_hdu'''

        RUNCF(self, cfitsio.ffcrhd)

    def create_image(self, axis_sizes, bitpix):
        '''Wrapper to ffcrim/fits_create_img'''

        # Is there a better way of making a ctypes array from a python tuple?
        num_dim = len(axis_sizes)
        naxes = (ctypes.c_long * num_dim)()
        for n in range(num_dim):
            naxes[n] = axis_sizes[n]

        RUNCF(self, cfitsio.ffcrim, bitpix, num_dim, naxes)

    def read_axes(self):

        #print '###### ', self.read_key('NAXIS')

        naxis = int(self.read_key('NAXIS'))

        naxes = (ctypes.c_long * naxis)()
        nfound = ctypes.c_int()
        RUNCF(self, cfitsio.ffgknj, "NAXIS", 1, naxis, naxes, ctypes.byref(nfound))

        return tuple(naxes)

    def get_bitpix(self):
        '''Wrapper to fits_get_img_type'''

        bitpix = ctypes.c_int()
        RUNCF(self, cfitsio.ffgidt, ctypes.byref(bitpix))
        return bitpix.value

    def read_key(self, keyname):
        '''Read value associated with keyword as string literal'''
        
        RUNCF(self, cfitsio.ffgkey, keyname, self.value, self.comment)
        return self.value.value

    def write_key(self, keyname, value, comment):
        '''Wrapper to ffuky/fits_update_key'''

        dtype = type(value)
        if dtype == int:
            value_ptr = ctypes.byref( ctypes.c_int(value) )
            cfitsio_type = cfitsio.TINT
        elif dtype == float:
            value_ptr = ctypes.byref( ctypes.c_double(value) )
            cfitsio_type = cfitsio.TDOUBLE
        elif dtype == str:
            value_ptr = value
            cfitsio_type = cfitsio.TSTRING
        else:
            raise ValueError('Key value must be primitive type')

        # Note this updates a key - ie overwrites if already exist
        RUNCF(self, cfitsio.ffuky, cfitsio_type, keyname, value_ptr, comment)

    def write_comment(self, comment):
        '''Wrapper to ffpcom/fits_write_comment'''
        RUNCF(self, cfitsio.ffpcom, comment)
        
    def write_history(self, history):
        '''Wrapper to ffphis/fits_write_history'''
        RUNCF(self, cfitsio.ffphis, history)
        
    def close(self):
        '''Wrapper to fits_close_file'''

        RUNCF(self, cfitsio.ffclos)



